from vec import Vec2
import math
import random
import pgzrun

WIDTH = 500
HEIGHT = 500

PEN_COLOR = (0, 200, 0)
BACK_COLOR = (0, 0, 0)
ptList = []
length = 400
pt_mid = Vec2(WIDTH / 2, HEIGHT / 2)
pt_a = Vec2(pt_mid.x - 100, pt_mid.y - 100)
pt_b = Vec2(pt_mid.x - 100, pt_mid.y + 100)
pt_c = Vec2(pt_mid.x + 100, pt_mid.y + 100)
pt_d = Vec2(pt_mid.x + 100, pt_mid.y - 100)
ptList.append(pt_a)
ptList.append(pt_b)
ptList.append(pt_c)
ptList.append(pt_d)

def distort(p):
    r = random.randint(1, 3)
    vx = random.randint(-100, 100)
    vy = random.randint(-100, 100)
    v = Vec2(vx, vy)
    v.normalize()
    dp = p + v * r

    return dp

def draw_fractal(pa, pb, n):
    if n == 1:
        if enableDistort == False:
            screen.draw.line(pa.raw_get(), pb.raw_get(), PEN_COLOR)
        else:
            screen.draw.line(distort(pa).raw_get(), distort(pb).raw_get(), PEN_COLOR)
    else:
        pts = []
        dir = pb - pa
        step = dir.length() / 4
        dir.normalize()
        pts.append(pa)

        pt1 = pa * 3 / 4 + pb / 4

        pts.append(pt1)

        vleftdir = Vec2(0, 0)
        vleftdir.x = -dir.y
        vleftdir.y = dir.x

        pt2 = pt1 + vleftdir * step

        pts.append(pt2)

        vupdir = Vec2(0, 0)
        vupdir.x = vleftdir.y
        vupdir.y = -vleftdir.x

        pt3 = pt2 + vupdir * step

        pts.append(pt3)

        vrightdir = Vec2(0, 0)
        vrightdir.x = vupdir.y
        vrightdir.y = -vupdir.x

        pt4 = pt3 + vrightdir * step

        pts.append(pt4)

        pt5 = pt4 + vrightdir * step

        pts.append(pt5)

        pt6 = pt5 + vupdir * step

        pts.append(pt6)

        pt7 = pt6 + vleftdir * step

        pts.append(pt7)

        pts.append(pb)

        for i in range(len(pts) - 1):
            draw_fractal(pts[i], pts[i + 1], n - 1)

count = len(ptList)
n = 2
enableDistort = False

def on_key_down(key):
    global enableDistort
    if key == keys.ESCAPE:
        quit() 
    elif key == keys.SPACE:
        enableDistort = not enableDistort

def update():
    pass

def draw():
    screen.fill(BACK_COLOR)
    
    for i in range(count):
        if i == count - 1:
            draw_fractal(ptList[count - 1], ptList[0], n)
        else:
            draw_fractal(ptList[i], ptList[i + 1], n) 

pgzrun.go()